import signal
import sys
import requests
from typing import List, Tuple, Dict
from sys import exit
import os
import argparse
import asyncio
from downloader import DownloadManager
import json

class ModelScope:
    """
    ModelScope 类，用于递归获取 ModelScope 仓库中的文件列表，并生成所需的元组列表。
    """

    def __init__(self, author : str, model_name : str, dataset_name : str, model_revision : str ='master', dataset_revision : str ='master'):
        """
        初始化 ModelScope 类实例。
        :param author: 模型作者
        :param model_name: 模型名称
        :param model_revision: 模型版本号 ,如 master
        :param dataset_name:数据集名称
        :param dataset_revision: 数据集版本号，如 master
        """
        self.author = author
        self.model_name = model_name
        self.model_revision = model_revision
        self.dataset_name = dataset_name
        self.dataset_revision = dataset_revision
        self.base_model_api_url = "https://modelscope.cn/api/v1/models"
        self.base_model_url = "https://www.modelscope.cn/models"
        self.base_dataset_api_url = "https://modelscope.cn/api/v1/datasets"
        self.base_dataset_url = "https://www.modelscope.cn/datasets"
         

    def get_model_files(self) -> List[Tuple[str, str, str]]:
        """
        递归获取model的所有文件列表。
        :return: 文件信息列表，每个元素为 (url, <author>/<model_name>/dir, file_name)
        """
        return self._fetch_model_files(dir_path="")

    def get_dataset_files(self) -> List[Tuple[str, str, str, int]]:
        """
        递归获取dataset的所有文件列表。
        :return: 文件信息列表，每个元素为 (url, <author>/<dataset_name>/dir, file_name)
        """
        dataset_id = self._fetch_dataset_id()
        return self._fetch_dataset_files(dataset_id=dataset_id, dir_path="")        

    def _fetch_model_files(self, dir_path : str ="") -> List[Tuple[str, str, str]]:
        """
        递归获取指定目录下的model所有文件信息。
        :param dir_path: 当前目录路径，默认为空表示根目录
        :return: 文件信息列表，每个元素为 (url, <author>/<model_name>/dir, file_name)
        """
        url = f"{self.base_model_api_url}/{self.author}/{self.model_name}/repo/files?Revision={self.model_revision}&Root={dir_path}"
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
        }
        response = requests.get(url, headers=headers)

        if response.status_code != 200:
            print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
            return []

        data = response.json()
        result = []

        if "Data" in data and "Files" in data["Data"]:
            for file_info in data["Data"]["Files"]:
                file_path = f"{dir_path}/{file_info['Name']}" if dir_path else file_info["Name"]
                if file_info["Type"] == "tree":  # 如果是目录，递归获取
                    result.extend(self._fetch_model_files(file_path))
                else:  # 如果是文件，生成元组并加入结果列表
                    file_name = file_info["Name"].split("/")[-1]  # 获取纯文件名
                    base_url = f"{self.base_model_url}/{self.author}/{self.model_name}/resolve/{self.model_revision}"
                    if dir_path:
                        full_url = f"{base_url}/{dir_path}/{file_name}"
                        relative_path = f"{self.author}/{self.model_name}/{dir_path}"
                    else:
                        full_url = f"{base_url}/{file_name}"
                        relative_path = f"{self.author}/{self.model_name}"
                    result.append((full_url, relative_path, file_name))

        return result

    def _fetch_dataset_id(self) -> str:
        url = f"{self.base_dataset_api_url}/{self.author}/{self.dataset_name}"
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
        }
        response = requests.get(url, headers=headers)        
        if response.status_code != 200:
            print(f"Failed to fetch dataset id from {url}. Status code: {response.status_code}")
            return ""
        data = response.json()
        if "Data" in data and "Id" in data["Data"]:
            return data['Data']['Id']
        else:
            return ""

    def _fetch_dataset_files(self, dataset_id, dir_path : str ="") -> List[Tuple[str, str, str, int]]:
        """
        递归获取指定目录下的dataset所有文件信息。
        :param dir_path: 当前目录路径，默认为空表示根目录
        :return: 文件信息列表，每个元素为 (url, <author>/<dataset_name>/dir, file_name)
        """
        url = f"{self.base_dataset_api_url}/{dataset_id}/repo/tree?Revision={self.dataset_revision}&Root={dir_path}&PageNumber=1&PageSize=2000	"
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3"
        }
        response = requests.get(url, headers=headers)
        if response.status_code != 200:
            print(f"Failed to fetch data from {url}. Status code: {response.status_code}")
            return []

        data = response.json()
        result = []

        if "Data" in data and "Files" in data["Data"]:
            for file_info in data["Data"]["Files"]:
                file_path = f"{dir_path}/{file_info['Name']}" if dir_path else file_info["Name"]
                if file_info["Type"] == "tree":  # 如果是目录，递归获取
                    result.extend(self._fetch_dataset_files(dataset_id, file_path))
                else:  # 如果是文件，生成元组并加入结果列表
                    file_name = file_info["Name"].split("/")[-1]  # 获取纯文件名
                    base_url = f"{self.base_dataset_url}/{self.author}/{self.dataset_name}/resolve/{self.dataset_revision}"
                    if dir_path:
                        full_url = f"{base_url}/{dir_path}/{file_name}"
                        relative_path = f"{self.author}/{self.dataset_name}/{dir_path}"
                    else:
                        full_url = f"{base_url}/{file_name}"
                        relative_path = f"{self.author}/{self.dataset_name}"
                    total_size = file_info['Size']
                    result.append((full_url, relative_path, file_name,total_size))

        return result

class Cmd:
    def __init__(self):
        """
        初始化命令行参数解析器
        """
        self.parser = argparse.ArgumentParser(
            description='ModelScope 工具',
            formatter_class=argparse.RawTextHelpFormatter
        )
        self.parser.add_argument(
            '-f', '--config', 
            type=str, 
            help='''配置文件路径。配置文件格式(JSON):
{
    "author": "作者名称",
    "model": "模型名称",
    "dataset": "数据集名称"
}
注意:
- author 为必填项
- model 和 dataset 至少需要填写一个'''
        )
        self.config = None
        self.manager = None  #  manager 属性,用于停止下载任务

    def _print_help(self):
        """
        打印帮助信息
        """
        self.parser.print_help()
        print("\n使用示例:")
        print("1. 使用配置文件:")
        print("   modelscope -f modelscope_config.json")
        print("\n配置文件示例(modelscope_config.json):")
        print('''{
    "author": "open-thoughts",
    "model": "",
    "dataset": "OpenThoughts-114k"
}''')

    def _load_config(self, config_file: str) -> Dict:
        """
        从配置文件加载参数        
        Args:
            config_file: 配置文件路径            
        Returns:
            Dict: 包含配置参数的字典
        """
        try:
            with open(config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                
                # 验证 author 字段
                author = config.get('author')
                if not author:
                    raise ValueError("配置文件中 'author' 字段不能为空")
                
                # 验证 model 和 dataset 字段
                model = config.get('model')
                dataset = config.get('dataset')
                if not model and not dataset:
                    raise ValueError("配置文件中 'model' 和 'dataset' 字段不能同时为空")
                    
                # 设置下载配置的默认值
                if 'download' not in config:
                    config['download'] = {}
                if 'max_concurrent_files' not in config['download']:
                    config['download']['max_concurrent_files'] = 5
                if 'max_concurrent_chunks' not in config['download']:
                    config['download']['max_concurrent_chunks'] = 3
                if 'chunk_size' not in config['download']:
                    config['download']['chunk_size'] = 50*1024*1024  # 50MB
                
                # 处理代理配置
                proxy_config = config['download'].get('proxy', {})
                # 默认不使用代理
                use_proxy = proxy_config.get('use_proxy', False)
                
                if use_proxy:
                    from downloader import ProxyConfig
                    proxy_settings = proxy_config.get('settings')
                    if not proxy_settings:
                        raise ValueError("启用代理时，必须提供 'settings' 配置")
                    
                    if isinstance(proxy_settings, str):
                        # 如果是字符串，直接作为URL
                        config['download']['proxy'] = ProxyConfig(url=proxy_settings)
                    elif isinstance(proxy_settings, dict):
                        # 如果是字典，可能包含更多配置
                        proxy_url = proxy_settings.get('http') or proxy_settings.get('https')
                        if not proxy_url:
                            raise ValueError("代理配置必须包含 'http' 或 'https' URL")
                        config['download']['proxy'] = ProxyConfig(
                            url=proxy_url,
                            auth=proxy_settings.get('auth'),
                            headers=proxy_settings.get('headers')
                        )
                else:
                    config['download']['proxy'] = None
                    
                return config
                
        except FileNotFoundError:
            print(f"错误: 找不到配置文件 {config_file}")
            return {}
        except json.JSONDecodeError:
            print(f"错误: 配置文件 {config_file} 格式不正确")
            return {}
        except ValueError as e:
            print(f"错误: {str(e)}")
            return {}
        except Exception as e:
            print(f"错误: 读取配置文件时发生错误: {str(e)}")
            return {}

    def check_options(self) -> bool:
        """
        检查命令行参数和配置文件
        Returns:
            bool: 检查是否通过
        """
        args = self.parser.parse_args()
        
        # 如果没有任何参数，显示帮助信息
        if not args.config:
            self._print_help()
            return False

        try:
            # 检查配置文件是否存在
            if not os.path.exists(args.config):
                print(f"错误: 找不到配置文件 {args.config}")
                self._print_help()
                return False

            # 如果指定了配置文件，从配置文件读取参数
            config = self._load_config(args.config)
            if not config:
                self._print_help()
                return False
            
            self.config = config
            return True
            
        except Exception as e:
            print(f"错误: {str(e)}")
            self._print_help()
            return False

    async def process(self, modelscope: ModelScope, manager: DownloadManager) -> None:
        """
        处理下载任务
        Args:
            modelscope: ModelScope类实例
            manager: DownloadManager类实例
        """
        try:
            author = self.config.get('author')
            model = self.config.get('model')
            dataset = self.config.get('dataset')
            self.manager = manager  # 保存 manager 实例            
            if model:
                # 获取作者特定模型的文件列表
                print(f"\n开始下载 ModelScope Model {author}/{model} 的文件 ...")
                files = modelscope.get_model_files()
                for url, directory, filename in files:
                    await manager.add_task(url, directory, filename)
                await manager.start_all()                
            elif dataset:
                # 获取作者特定数据集的文件列表
                print(f"\n开始下载 ModelScope Dataset {author}/{dataset} 的文件 ...")
                files = modelscope.get_dataset_files()
                for url, directory, filename, total_size in files:
                    await manager.add_task(url, directory, filename, total_size)
                await manager.start_all()
            else:
                self._print_help()
                return
            
        except Exception as e:
            print(f"错误: {str(e)}")
            self._print_help()
            return

    async def handle_signal(self):
        """处理信号"""
        if self.manager:
            await self.manager.cancel_all_tasks()

    def get_config(self):
        return self.config

# 修改主函数
async def main():
    """主处理程序"""
    # 命令行处理类
    cmd = Cmd()
    
    # 检查命令行参数
    if not cmd.check_options():
        exit(1)
    
    # 获取配置
    config = cmd.get_config()
    
    # 初始化实例
    modelscope = ModelScope(author=config['author'], model_name=config['model'], 
                dataset_name=config['dataset'], model_revision=config['model_revision'], 
                dataset_revision=config['dataset_revision'])

    # 配置下载管理器
    manager = DownloadManager(
        max_concurrent_files=config['download']['max_concurrent_files'],
        max_concurrent_chunks=config['download']['max_concurrent_chunks'],
        chunk_size=config['download']['chunk_size'],
        default_proxy=config['download']['proxy'],  # 从配置文件读取代理设置
        hf_token=None
    )

    # 添加基本的信号处理
    def signal_handler(sig, frame):
        print('\n收到中断信号，正在退出程序...')
        # 创建一个任务来处理取消操作
        loop = asyncio.get_running_loop()
        loop.create_task(cmd.handle_signal())
    
    # 设置信号处理器
    signal.signal(signal.SIGINT, signal_handler)
    
    try:
        # 执行下载任务
        await cmd.process(modelscope, manager)
    finally:
        # 恢复默认信号处理
        signal.signal(signal.SIGINT, signal.default_int_handler)    


# 主函数
if __name__ == "__main__":
    # Windows平台需要使用不同的事件循环策略
    if os.name == 'nt':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
    
    try:
        # 运行主程序
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序已退出")
    except SystemExit:
        pass